package my.tree;

import java.util.Comparator;

/**
 * <p>
 * Red Black Tree: 红黑树
 * </p>
 *
 * @author AJun
 * @since 2020/8/29
 */
@SuppressWarnings("UnusedReturnValue")
public class RBTree<E> extends BBST<E> {

    /**
     * 红色节点: false
     */
    private static final boolean RED = false;

    /**
     * 黑色节点: true
     */
    private static final boolean BLACK = true;

    /**
     * 空参构造
     * <p>
     * 没有传入比较器，则存储的元素必须 implements Comparable 接口实现比较逻辑，否则会报异常
     * </p>
     */
    public RBTree() {
        this(null);
    }

    /**
     * 添加之后的操作
     *
     * @param node 添加的节点
     */
    @Override
    protected void afterAdd(Node<E> node) {
        // 父节点
        Node<E> parent = node.parent;

        // 添加的是根节点或者上溢到达了根节点，则直接将该节点染成黑色然后直接添加即可
        if (parent == null) {
            black(node);
            return;
        }

        // 如果父节点为黑色，则无需调整，直接返回
        if (isBlack(parent)) {
            return;
        }

        // 叔父节点
        Node<E> uncle = parent.sibling();
        // 祖父节点(无论是否上溢，都需要将其染成红色)
        Node<E> grand = red(parent.parent);

        // 上溢
        if (isRed(uncle)) {
            // 叔父节点是红色，上溢，将父节点和叔父节点染成黑色
            black(parent);
            black(uncle);
            // 将祖父节点染成红色当作新添加的节点向上合并
            afterAdd(grand);
            return;
        }

        // 非上溢: 叔父节点不是红色，染色并旋转
        if (parent.isLeftChild()) { // L
            if (node.isLeftChild()) { // LL
                // 父节点染成黑色，祖父节点染成红色，祖父节点右旋转
                black(parent);
            } else { // LR
                // 当前节点染成黑色，祖父节点染成红色，父节点左旋转，祖父节点右旋转
                black(node);
                rotateLeft(parent);
            }
            // 祖父节点进行右旋转
            rotateRight(grand);
        } else { // R
            if (node.isLeftChild()) { // RL
                // 当前节点染成黑色，祖父节点染成红色，父节点右旋转，祖父节点左旋转
                black(node);
                rotateRight(parent);
            } else { // RR
                // 父节点染成黑色，祖父节点染成红色，祖父节点进行左旋转
                black(parent);
            }
            // 祖父节点进行左旋转
            rotateLeft(grand);
        }
    }

    /**
     * 删除之后的操作
     *
     * @param node 被删除的节点 或者 用于取代被删除节点的子节点(当被删除节点度为1时)
     */
    @Override
    protected void afterRemove(Node<E> node) {
        // 被删除节点是红色 或者 用于取代被删除节点的子节点是红色
        if (isRed(node)) {
            black(node);
            return;
        }

        // 被删除节点的父节点
        Node<E> parent = node.parent;

        // 被删除节点是根节点
        if (parent == null) {
            return;
        }

        // 被删除节点是黑色叶子节点(下溢)
        // 判断被删除节点是左边还是右边
        boolean left = (parent.left == null) || (node.isLeftChild());
        // 被删除节点的兄弟节点
        Node<E> sibling = (left ? parent.right : parent.left);
        if (left) {
            // 被删除节点在左边，兄弟节点在右边
            if (isRed(sibling)) {
                // 兄弟节点是红色
                // 兄弟节点染成黑色，父节点染成红色
                black(sibling);
                red(parent);
                // 将父节点进行左旋转
                rotateLeft(parent);
                // 更换兄弟节点
                sibling = parent.right;
            }
            // 兄弟节点是黑色
            if (isBlack(sibling.left) && isBlack(sibling.right)) {
                // 兄弟节点左右都是BLACK节点，即: 兄弟节点没有一个红色子节点，父节点需要向下与兄弟节点合并
                // 判断父节点是不是黑色(父节点是红色则不会下溢，否则会下溢)
                boolean parentBlack = isBlack(parent);
                // 父节点染成黑色，兄弟节点染成红色
                black(parent);
                red(sibling);

                // 父节点下溢
                if (parentBlack) {
                    afterRemove(parent);
                }
            } else {
                // 兄弟节点至少有一个红色子节点，需要向兄弟节点借元素
                if (isBlack(sibling.right)) {
                    // 兄弟节点只有一个红色子节点且红色子节点是右子节点
                    // 对兄弟节点进行左旋转
                    rotateRight(sibling);

                    // 更换兄弟节点，此时兄弟节点由于旋转已经被更改
                    sibling = parent.right;
                }

                // 染色: 旋转之后的中心节点继承父节点颜色，左右节点染成黑色
                // 先染色: 兄弟节点继承父节点颜色，父节点和兄弟节点的子节点染成黑色
                color(sibling, colorOf(parent));
                black(sibling.right);
                black(parent);

                // 兄弟节点只有一个红色子节点且红色子节点是左子节点 或 兄弟节点有俩个红色个节点
                // 都是对父节点进行左旋转
                rotateLeft(parent);
            }
        } else {
            // 被删除节点在右边，兄弟节点在左边
            if (isRed(sibling)) {
                // 兄弟节点是红色
                // 兄弟节点染成黑色，父节点染成红色
                black(sibling);
                red(parent);
                // 将父节点进行右旋转
                rotateRight(parent);
                // 更换兄弟节点
                sibling = parent.left;
            }
            // 兄弟节点是黑色
            if (isBlack(sibling.left) && isBlack(sibling.right)) {
                // 兄弟节点左右都是BLACK节点，即: 兄弟节点没有一个红色子节点，父节点需要向下与兄弟节点合并
                // 判断父节点是不是黑色(父节点是红色则不会下溢，否则会下溢)
                boolean parentBlack = isBlack(parent);
                // 父节点染成黑色，兄弟节点染成红色
                black(parent);
                red(sibling);

                // 父节点下溢
                if (parentBlack) {
                    afterRemove(parent);
                }
            } else {
                // 兄弟节点至少有一个红色子节点，需要向兄弟节点借元素
                if (isBlack(sibling.left)) {
                    // 兄弟节点只有一个红色子节点且红色子节点是右子节点
                    // 对兄弟节点进行左旋转
                    rotateLeft(sibling);

                    // 更换兄弟节点，此时兄弟节点由于旋转已经被更改
                    sibling = parent.left;
                }

                // 染色: 旋转之后的中心节点继承父节点颜色，左右节点染成黑色
                // 先染色: 兄弟节点继承父节点颜色，父节点和兄弟节点的子节点染成黑色
                color(sibling, colorOf(parent));
                black(sibling.left);
                black(parent);

                // 兄弟节点只有一个红色子节点且红色子节点是左子节点 或 兄弟节点有俩个红色个节点
                // 都是对父节点进行右旋转
                rotateRight(parent);
            }
        }
    }

    /**
     * 有参构造
     *
     * @param comparator 用于元素比较的比较器
     */
    public RBTree(Comparator<E> comparator) {
        super(comparator);
    }

    /**
     * 将传入节点染色
     *
     * @param node  需要染色的节点
     * @param color 该节点需要染的颜色
     * @return 染色后的节点
     */
    private Node<E> color(Node<E> node, boolean color) {
        if (node == null) {
            return null;
        }

        ((RBNode<E>) node).color = color;
        return node;
    }

    /**
     * 将传入节点染成红色
     *
     * @param node 需要染成红色的节点
     * @return 染色后的节点
     */
    private Node<E> red(Node<E> node) {
        return color(node, RED);
    }

    /**
     * 将传入节点染成黑色
     *
     * @param node 需要染成黑色的节点
     * @return 染色后的节点
     */
    private Node<E> black(Node<E> node) {
        return color(node, BLACK);
    }

    /**
     * 获取传入节点的颜色
     *
     * @param node 红黑树节点
     * @return 该节点颜色: true: 黑色; false: 红色
     */
    private boolean colorOf(Node<E> node) {
        // 如果为空则认为是黑色
        return node == null ? BLACK : ((RBNode<E>) node).color;
    }

    /**
     * 判断该节点是否为红色
     *
     * @param node 红黑树节点
     * @return 是否为红色: true: 红色; false: 不是黑色
     */
    private boolean isRed(Node<E> node) {
        return colorOf(node) == RED;
    }

    /**
     * 判断该节点是否为黑色
     *
     * @param node 红黑树节点
     * @return 是否为黑色: true: 黑色; false: 不是黑色
     */
    private boolean isBlack(Node<E> node) {
        return colorOf(node) == BLACK;
    }

    /**
     * 创建红黑树节点
     *
     * @param element 节点的元素内容
     * @param parent  节点的父节点
     * @return 红黑树节点
     */
    @Override
    protected Node<E> createNode(E element, Node<E> parent) {
        return new RBNode<>(element, parent);
    }

    /**
     * 红黑树节点
     */
    private static class RBNode<E> extends Node<E> {
        /**
         * 添加的红黑树节点默认为红色，可以保证除第四条性质外的性质都满足
         */
        boolean color = RED;

        public RBNode(E element, Node<E> parent) {
            super(element, parent);
        }

        @Override
        public String toString() {
            String res = "";
            if (color == RED) {
                res = "R_";
            }
            return res + element.toString();
        }

    }

}
